import asyncio
import websockets
import json
import base64
import _thread as thread
from sse_starlette.sse import EventSourceResponse
from typing import AsyncGenerator
from dataclasses import dataclass
from websocket import WebSocketApp
import websocket
import hashlib
import hmac
from urllib.parse import urlencode
from datetime import datetime
from time import mktime
from wsgiref.handlers import format_date_time
import ssl
import time

STATUS_FIRST_FRAME = 0
STATUS_CONTINUE_FRAME = 1
STATUS_LAST_FRAME = 2

@dataclass
class ASRConfig:
    app_id: str
    api_key: str
    api_secret: str
    host: str = "ws-api.xfyun.cn"
    path: str = "/v2/iat"

class SparkASRService:
    def __init__(self, config: ASRConfig):
        self.config = config
        self.clients = {}  # client_id -> websocket
        self.audio_buffers = {}  # client_id -> audio buffer
        self.asr_connections = {}  # client_id -> ASR WebSocket
        
    async def handle_ws_client(self, websocket, path):
        client_id = id(websocket)
        self.clients[client_id] = websocket
        self.audio_buffers[client_id] = bytearray()
        last_data_time = asyncio.get_event_loop().time()
        
        try:
            while True:
                try:
                    # Wait for message with timeout
                    message = await asyncio.wait_for(
                        websocket.recv(),
                        timeout=0.5  # 500ms timeout
                    )
                    
                    if isinstance(message, bytes):
                        self.audio_buffers[client_id].extend(message)
                        last_data_time = asyncio.get_event_loop().time()
                    
                except asyncio.TimeoutError:
                    # If no data for 500ms and we have buffered data, process it
                    if (asyncio.get_event_loop().time() - last_data_time >= 0.5 and 
                        len(self.audio_buffers[client_id]) > 0):
                        await self.process_audio(client_id)
                        last_data_time = asyncio.get_event_loop().time()
                    continue
                    
                except websockets.exceptions.ConnectionClosed:
                    break
                    
        finally:
            # Clean up resources
            if client_id in self.clients:
                del self.clients[client_id]
            del self.audio_buffers[client_id]
            if client_id in self.asr_connections:
                self.asr_connections[client_id].close()
                del self.asr_connections[client_id]

    def _create_url(self):
        url = 'wss://ws-api.xfyun.cn/v2/iat'
        # url = 'wss://iat-api.xfyun.cn/v2/iat'
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = "host: " + "ws-api.xfyun.cn" + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + "/v2/iat " + "HTTP/1.1"
        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(self.config.api_secret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = "api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"" % (
            self.config.api_key, "hmac-sha256", "host date request-line", signature_sha)
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": "ws-api.xfyun.cn"
        }
        # 拼接鉴权参数，生成url
        url = url + '?' + urlencode(v)
        # print("date: ",date)
        # print("v: ",v)
        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，比对相同参数时生成的url与自己代码生成的url是否一致
        # print('websocket url :', url)
        return url
    def create_asr_connection(self, client_id: int):
        """Create WebSocket connection to ASR service"""
        def on_message(ws, message):
            try:
                data = json.loads(message)
                if data["code"] != 0:
                    print(f"ASR error: {data['message']}")
                    return
                
                result = ""
                for item in data["data"]["result"]["ws"]:
                    for w in item["cw"]:
                        result += w["w"]
                print(f"ASR result: {result}")
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                loop.run_until_complete(self.send_sse_event(
                    client_id,
                    {"text": result, "is_final": data["data"]["status"] == 2}
                ))
                loop.close()
            except Exception as e:
                print(f"Error processing ASR message: {e}")

        def on_error(ws, error):
            print(f"ASR connection error: {error}")

        def on_close(ws, *args):
            print("ASR connection closed")
            if client_id in self.asr_connections:
                del self.asr_connections[client_id]

        def on_open(ws):
            intervel = 0.04  # 发送音频间隔(单位:s)
            def run(*args):
                # Send first frame with config
                d = {
                    "common": {"app_id": self.config.app_id},
                    "business": {
                        "domain": "iat",
                        "language": "zh_cn",
                        "accent": "mandarin",
                        "vinfo": 1,
                        "vad_eos": 3000
                    },
                    "data": {
                        "status": STATUS_FIRST_FRAME,
                        "format": "audio/L16;rate=16000",
                        "encoding": "raw"
                    }
                }
                ws.send(json.dumps(d))
                print("send first frame")
                
                # Continuously send audio data
                last_send_time = time.time()
                while True:
                    try:
                        if client_id in self.audio_buffers:
                            buffer = self.audio_buffers[client_id]
                            print("buffer size: ", len(buffer))
                            if len(buffer) >= 8000:
                                chunk = buffer[:8000]
                                del buffer[:8000]
                                
                                status = STATUS_CONTINUE_FRAME
                                if len(buffer) == 0:
                                    status = STATUS_LAST_FRAME
                                
                                data = {
                                    "data": {
                                        "status": status,
                                        "format": "audio/L16;rate=16000",
                                        "audio": base64.b64encode(chunk).decode(),
                                        "encoding": "raw"
                                    }
                                }
                                ws.send(json.dumps(data))
                                last_send_time = time.time()
                                print("send...")
                            
                            # Timeout if no data for 500ms
                            if time.time() - last_send_time > 0.4 and len(buffer) > 0:
                                print(f"timeout last frame:{len(buffer)}")
                                chunk = buffer[:]
                                del buffer[:]
                                
                                data = {
                                    "data": {
                                        "status": STATUS_LAST_FRAME,
                                        "format": "audio/L16;rate=16000",
                                        "audio": base64.b64encode(chunk).decode(),
                                        "encoding": "raw"
                                    }
                                }
                                ws.send(json.dumps(data))
                                break
                            elif len(buffer) <=0:
                                break
                            
                        time.sleep(intervel)
                    except Exception as e:
                        print(f"Error in ASR send loop: {e}")
                        break

            thread.start_new_thread(run, ())

        # Generate auth URL
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        signature_origin = f"host: {self.config.host}\ndate: {date}\nGET {self.config.path} HTTP/1.1"
        signature_sha = hmac.new(
            self.config.api_secret.encode(),
            signature_origin.encode(),
            digestmod=hashlib.sha256
        ).digest()
        signature = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization = f'api_key="{self.config.api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature}"'
        authorization = base64.b64encode(authorization.encode('utf-8')).decode(encoding='utf-8')

        url = f"wss://{self.config.host}{self.config.path}?authorization={authorization}&date={date}&host={self.config.host}"

        # Create headers with required authentication parameters
        headers = {
            "Authorization": authorization,
            "Date": date,
            "Host": self.config.host
        }
        url = self._create_url()
        ws = websocket.WebSocketApp(
            url,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close,
            on_open=on_open
        )
        self.asr_connections[client_id] = ws
        return ws

    async def process_audio(self, client_id: int):
        """Process audio buffer and send to ASR"""
        buffer = self.audio_buffers[client_id]
        if len(buffer) < 8000:  # Wait for enough data
            return
            
        chunk = buffer[:8000]
        del buffer[:8000]
        
        if client_id not in self.asr_connections:
            self.asr_connections[client_id] = self.create_asr_connection(client_id)
            thread.start_new_thread(
                self.asr_connections[client_id].run_forever,
                (None,{"cert_reqs": ssl.CERT_NONE})
            )
            await asyncio.sleep(0.1)
        
        # Get connection and check status
        ws = self.asr_connections.get(client_id)
        if not ws or not ws.sock or not ws.sock.connected:
            if client_id in self.asr_connections:
                del self.asr_connections[client_id]
            return
        # Determine frame status
        status = STATUS_CONTINUE_FRAME
        if len(self.audio_buffers[client_id]) == 0:
            status = STATUS_LAST_FRAME
            
        # Send audio chunk to ASR
        data = {
            "data": {
                "status": status,
                "format": "audio/L16;rate=16000",
                "audio": base64.b64encode(chunk).decode(),
                "encoding": "raw"
            }
        }
        self.asr_connections[client_id].send(json.dumps(data))
        
        if status == STATUS_LAST_FRAME:
            # Let ASR connection close naturally after receiving last frame
            del self.asr_connections[client_id]
    
    async def send_sse_event(self, client_id: int, data: dict):
        """Send SSE event to client"""
        if client_id in self.clients:
            try:
                await self.clients[client_id].send(json.dumps(data))
            except Exception as e:
                print(f"Error sending to client {client_id}: {e}")
    
    async def run_server(self, host: str = "0.0.0.0", port: int = 8765):
        """Start WebSocket server"""
        async with websockets.serve(
            self.handle_ws_client, host, port
        ):
            print(f"ASR WebSocket server started at ws://{host}:{port}")
            await asyncio.Future()  # run forever

if __name__ == "__main__":
    config = ASRConfig(
        app_id="f4b3bfd8",
        api_key="3f6b034a6397fab4d4478345f42340e2",
        api_secret="MDdhM2I3NjM0NjViOTVlY2JkZjQ5NTZj"
    )
    service = SparkASRService(config)
    asyncio.run(service.run_server())
